import argparse
import auxil
import time
from hyper_pytorch import *
import scipy.io as sio
import torch
import torch.nn.parallel
from torchvision.transforms import *

import models.resnet as RN
import models.presnet as PYRM


#
#
# def load_hyper(args):
#
# 	return train_loader, test_loader, numberofclass, bands


def train(trainloader, model, criterion, optimizer, epoch, use_cuda):
	model.train()
	accs   = np.ones((len(trainloader))) * -1000.0
	losses = np.ones((len(trainloader))) * -1000.0
	for batch_idx, (inputs, targets) in enumerate(trainloader):
		if use_cuda:
			inputs, targets = inputs.cuda(), targets.cuda()
		inputs, targets = torch.autograd.Variable(inputs), torch.autograd.Variable(targets)
		outputs = model(inputs)
		loss = criterion(outputs, targets)
		losses[batch_idx] = loss.item()
		accs[batch_idx] = auxil.accuracy(outputs.data, targets.data)[0].item()
		optimizer.zero_grad()
		loss.backward()
		optimizer.step()
	return (np.average(losses), np.average(accs))


def test(testloader, model, criterion, epoch, use_cuda):
	model.eval()
	accs   = np.ones((len(testloader))) * -1000.0
	losses = np.ones((len(testloader))) * -1000.0
	for batch_idx, (inputs, targets) in enumerate(testloader):
		if use_cuda:
			inputs, targets = inputs.cuda(), targets.cuda()
		inputs, targets = torch.autograd.Variable(inputs, volatile=True), torch.autograd.Variable(targets)
		outputs = model(inputs)
		losses[batch_idx] = criterion(outputs, targets).item()
		accs[batch_idx] = auxil.accuracy(outputs.data, targets.data, topk=(1,))[0].item()
	return (np.average(losses), np.average(accs))




def predict(testloader, model, criterion, use_cuda):
	model.eval()
	predicted = []
	for batch_idx, (inputs, targets) in enumerate(testloader):
		if use_cuda: inputs = inputs.cuda()
		inputs, targets = torch.autograd.Variable(inputs, volatile=True), torch.autograd.Variable(targets)
		[predicted.append(a) for a in model(inputs).data.cpu().numpy()] 
	return np.array(predicted)


def adjust_learning_rate(optimizer, epoch, args):
	lr = args.lr * (0.1 ** (epoch // 150)) * (0.1 ** (epoch // 225))
	for param_group in optimizer.param_groups:
		param_group['lr'] = lr


def main():
	parser = argparse.ArgumentParser(description='PyTorch DCNNs Training')
	parser.add_argument('--epochs', default=200, type=int, help='number of total epochs to run')
	parser.add_argument('--lr', '--learning-rate', default=0.1, type=float, help='initial learning rate')
	parser.add_argument('--components', default=None, type=int, help='dimensionality reduction')
	parser.add_argument('--dataset', default=None, type=str, help='dataset (options: IP, UH, PU, KSC)')
	parser.add_argument('--tr_percent', default=0.15, type=float, help='samples of train set')
	parser.add_argument('--tr_bsize', default=100, type=int, help='mini-batch train size (default: 100)')
	parser.add_argument('--te_bsize', default=1000, type=int, help='mini-batch test size (default: 1000)')
	parser.add_argument('--depth', default=32, type=int, help='depth of the network (default: 32)')
	parser.add_argument('--alpha', default=50, type=int, help='number of new channel increases per depth (default: 12)')
	parser.add_argument('--inplanes', dest='inplanes', default=16, type=int, help='bands before blocks')
	parser.add_argument('--no-bottleneck', dest='bottleneck', action='store_false', help='to use basicblock (default: bottleneck)')
	parser.add_argument('--spatialsize', dest='spatialsize', default=11, type=int, help='spatial-spectral patch dimension')
	parser.add_argument('--momentum', default=0.9, type=float, help='momentum')
	parser.add_argument('--weight-decay', '--wd', default=1e-4, type=float, help='weight decay (default: 1e-4)')

	parser.set_defaults(bottleneck=True)
	best_err1 = 100

	args = parser.parse_args()
	state = {k: v for k, v in args._get_kwargs()}

	# 导入数据集，创建数据矩阵

	data, labels, numclass = auxil.loadData(args.dataset, num_components=args.components)
	pixels, labels = auxil.createImageCubes(data, labels, windowSize=args.spatialsize, removeZeroLabels=True)
	n_bands = pixels.shape[-1];
	num_classes = len(np.unique(labels))

	print('数据集',args.dataset,'数据集尺寸',data.shape,'类别数',num_classes)

	# 10次重复试验

	Experiment_results = np.zeros([num_classes + 5, 12])

	for iter in range(1):

		x_train, x_test, y_train, y_test = auxil.split_data(pixels, labels, args.tr_percent, args.dataset,
															splitdset="custom", rand_state=iter)
		train_hyper = HyperData((np.transpose(x_train, (0, 3, 1, 2)).astype("float32"), y_train), None)
		test_hyper = HyperData((np.transpose(x_test, (0, 3, 1, 2)).astype("float32"), y_test), None)
		kwargs = {'num_workers': 1, 'pin_memory': True}
		train_loader = torch.utils.data.DataLoader(train_hyper, batch_size=args.tr_bsize, shuffle=True, **kwargs)
		test_loader = torch.utils.data.DataLoader(test_hyper, batch_size=args.te_bsize, shuffle=False, **kwargs)

		#train_loader, test_loader, num_classes, n_bands = load_hyper(args)

		# Use CUDA
		use_cuda = torch.cuda.is_available()
		if use_cuda: torch.backends.cudnn.benchmark = True

		if args.spatialsize < 9: avgpoosize = 1
		elif args.spatialsize <= 11: avgpoosize = 2
		elif args.spatialsize == 15: avgpoosize = 3
		elif args.spatialsize == 19: avgpoosize = 4
		elif args.spatialsize == 21: avgpoosize = 5
		elif args.spatialsize == 27: avgpoosize = 6
		elif args.spatialsize == 29: avgpoosize = 7
		else: print("spatialsize no tested")

		model = PYRM.pResNet(args.depth, args.alpha, num_classes, n_bands, avgpoosize, args.inplanes, bottleneck=args.bottleneck) # for PyramidNet
		model = torch.nn.DataParallel(model)
		if use_cuda: model = model.cuda()

		criterion = torch.nn.CrossEntropyLoss()
		#optimizer = torch.optim.Adam(model.parameters())
		optimizer = torch.optim.SGD(model.parameters(), args.lr,
									momentum=args.momentum,
									weight_decay=args.weight_decay, nesterov=True)


		title = 'HYPER-' + args.dataset

		best_acc = -1

		time1=time.time()# trn time
		for epoch in range(args.epochs):
			adjust_learning_rate(optimizer, epoch, args)

			train_loss, train_acc = train(train_loader, model, criterion, optimizer, epoch, use_cuda)
			#test_loss, test_acc = test(test_loader, model, criterion, epoch, use_cuda)

			print("EPOCH", epoch, "TRAIN LOSS", train_loss, "TRAIN ACCURACY", train_acc, end=',')
			print(' ')
			#print("LOSS", test_loss, "ACCURACY", test_acc)
			# save model
		time2 = time.time()

		time3 = time.time()  # tes time
		test_loss, test_acc = test(test_loader, model, criterion, epoch, use_cuda)
		time4 = time.time()

		print("FINAL:      LOSS", test_loss, "ACCURACY", test_acc)

		if test_acc > best_acc:
			state = {
					'epoch': epoch + 1,
					'state_dict': model.state_dict(),
					'acc': test_acc,
					'best_acc': best_acc,
					'optimizer' : optimizer.state_dict(),
			}
			torch.save(state, "best_model_"+str(args.dataset)+".pth.tar")
			best_acc = test_acc


		checkpoint = torch.load("best_model_"+str(args.dataset)+".pth.tar")
		best_acc = checkpoint['best_acc']
		start_epoch = checkpoint['epoch']
		model.load_state_dict(checkpoint['state_dict'])
		optimizer.load_state_dict(checkpoint['optimizer'])

		y_pred=np.argmax(predict(test_loader, model, criterion, use_cuda), axis=1)

		classification, confusion, results = auxil.reports(y_pred, np.array(test_loader.dataset.__labels__()), args.dataset)
		print('Experiments {} times'.format(iter),args.dataset, results)

		Experiment_results[:3,iter]=results[:3]
		Experiment_results[3,iter]=time2-time1
		Experiment_results[4, iter]=time4-time3
		Experiment_results[5:,iter]=results[3:]

		Experiment_results[:,-2] = np.mean(Experiment_results[:,:iter+1], axis=1)
		Experiment_results[:,-1] = np.std(Experiment_results[:, :iter+1], axis=1)

		sio.savemat('DPRN_'+str(args.dataset)+'.mat',{'data':Experiment_results})

	print('Training & Testing finished!')

if __name__ == '__main__':
	main()

